from utils import get_db_connection
import psycopg2.sql as sql
from mapping_table_operations import MappingTableOperations

class OriginalTableOperations:
    def __init__(self, config):
        self.config = config
        self.connection = get_db_connection(config)
        self.mapping_ops = MappingTableOperations(config)

    def create_sensitive_data_table(self, table_name, columns):
        self.mapping_ops.create_sensitive_data_table(table_name, columns)

    def insert_sensitive_data(self, table_name, sensitive_columns, new_table_name):
        self.mapping_ops.insert_sensitive_data(table_name, sensitive_columns, new_table_name)

    def update_original_table(self, table_name, sensitive_columns, new_table_name):
        with self.connection.cursor() as cursor:
            for column in sensitive_columns:
                if self.config.db_type == "postgresql":
                    # 使用 JOIN 更新原表中的敏感列数据为映射表中的 id
                    cursor.execute(sql.SQL("""
                        UPDATE {}
                        SET {} = {}.id
                        FROM {}
                        WHERE {}.{} = {}.{}
                    """).format(
                        sql.Identifier(table_name),
                        sql.Identifier(column),
                        sql.Identifier(new_table_name),
                        sql.Identifier(new_table_name),
                        sql.Identifier(table_name),
                        sql.Identifier(column),
                        sql.Identifier(new_table_name),
                        sql.Identifier(column)
                    ))
                elif self.config.db_type == "mysql":
                    # 使用 JOIN 更新原表中的敏感列数据为映射表中的 id
                    cursor.execute(sql.SQL("""
                        UPDATE {}
                        JOIN {} ON {}.{} = {}.{}
                        SET {}.{} = {}.id
                    """).format(
                        sql.Identifier(table_name),
                        sql.Identifier(new_table_name),
                        sql.Identifier(table_name),
                        sql.Identifier(column),
                        sql.Identifier(new_table_name),
                        sql.Identifier(column),
                        sql.Identifier(table_name),
                        sql.Identifier(column),
                        sql.Identifier(new_table_name)
                    ))
                self.connection.commit()

    def close_connection(self):
        self.connection.close()